#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <stdio.h>
#include "buddy.h"
#include "util.h"
#include "mm.h"

typedef struct buddy_t
{
	int pool_size;

	/* min alloc unit size */
	int min_size;

	/* allocate max & min order */
	int min_order;
	int max_order;

	/* buddy index */
	int8_t* tree;
	int tree_size;

	/* memory pool */
	char* pool;
}buddy_t;

#define LMA_BUDDY_UNUSED 0
#define LMA_BUDDY_USED 1
#define LMA_BUDDY_SPLIT 2
#define LMA_BUDDY_FULL 3

/* get memory size by index */
#define LMA_BUDDY_SIZE_FROM_INDEX(buddy, index) \
	(buddy->pool_size >> ILOG2(index))

/* get memory shift by index */
#define LMA_BUDDY_SHIFT_FROM_INDEX(buddy, index) \
	(buddy->pool_size >> ILOG2(index )) * (index - (1 << ILOG2(index))) 

/* get memory by index */
#define LMA_BUDDY_MEM_FROM_INDEX(buddy, index) \
	buddy->pool + LMA_BUDDY_SHIFT_FROM_INDEX(buddy, index)

struct buddy_t* buddy_init(size_t size, size_t min_alloc_size)
{
	/* round up size by 2^n */
	size_t real_size = ROUNDUP_POWOF_2(size);
	size_t real_min_size = ROUNDUP_POWOF_2(min_alloc_size);
	if(real_size > LMA_MAX_SIZE || real_size < LMA_MIN_SIZE
		|| real_min_size > LMA_MAX_SIZE || real_min_size < LMA_MIN_SIZE
		|| real_min_size >= real_size) 
	{
		return NULL;
	}

	/* alloc struct */
	struct buddy_t* buddy = (struct buddy_t*)MALLOC(sizeof(buddy_t));
	buddy->pool_size = real_size;
	buddy->min_size = real_min_size;
	buddy->max_order = ILOG2(buddy->pool_size) ;
	buddy->min_order = ILOG2(buddy->min_size);

	/* alloc buddy index */
	buddy->tree_size = buddy->pool_size * 2 / buddy->min_size;
	buddy->tree = (int8_t*)MALLOC(sizeof(int8_t) * buddy->tree_size);
	int index;
	for(index = 0; index < buddy->tree_size; index ++)
		buddy->tree[index] = LMA_BUDDY_UNUSED;

	/* alloc memory pool */
	buddy->pool = MALLOC(buddy->pool_size);
	return buddy;
}

int buddy_release(struct buddy_t* buddy)
{
	if(buddy)
	{
		if(!buddy->tree || !buddy->pool)
			assert(0);

		if(buddy->tree)
			FREE(buddy->tree);
		if(buddy->pool)
			FREE(buddy->pool);
		
		FREE(buddy);
		buddy = 0;
	}
	return 0;
}

void* buddy_realloc(struct buddy_t* buddy, void* mem, size_t nbytes)
{
	if(!buddy)
		return NULL;
	if(!mem)
		return buddy_alloc(buddy, nbytes);
	
	if(nbytes == 0)
	{
		buddy_free(buddy, mem);
		return NULL;
	}

	/* calculate original size */
	#ifdef __x86_64__
	int offset = (int)((uint64_t)mem - (uint64_t)buddy->pool);
#else
	int offset = (int)((uint32_t)mem - (uint32_t)buddy->pool);
#endif
	assert(offset <= buddy->pool_size);
	int index = 1;
	int step = buddy->pool_size;
	int current = 0;
	size_t mem_size = 0;
	while(1)
	{	
		if(buddy->tree[index] == LMA_BUDDY_USED)
		{
			mem_size = LMA_BUDDY_SIZE_FROM_INDEX(buddy, index);
			break;
		}
		/* find dest by dichotomy */
		step/= 2;
		index *= 2;
		if(offset >= current + step)
		{
			current += step;
			index ++;
		}
	}

	/* alloc new memory */
	void* new_mem = buddy_alloc(buddy, nbytes);
	if(new_mem)
	{
		memcpy(new_mem, mem, (mem_size < nbytes ? mem_size : nbytes));
	}
	buddy_free(buddy, mem);
	return new_mem;
}

void* buddy_alloc(struct buddy_t* buddy, size_t nbytes)
{
	if(nbytes < 0 || !buddy)
	{
		return NULL;
	}

	/* calculate malloc size */
	size_t malloc_size = ROUNDUP_POWOF_2(nbytes);
	if(nbytes > buddy->pool_size)
	{
		return NULL;
	}
	if(malloc_size < buddy->min_size)
		malloc_size = buddy->min_size;

#ifndef LMA_ALLOC_CHECK
#define LMA_ALLOC_CHECK(index) \
	if(index == 1) \
	{ \
		return NULL; \
	}
	
	/* depth first search */
	int index = 1;
	while(index < buddy->tree_size)
	{
		if(buddy->tree[index] == LMA_BUDDY_UNUSED)
		{
			size_t mem_size = LMA_BUDDY_SIZE_FROM_INDEX(buddy, index);

			/* mark used */
			if(mem_size == malloc_size)
			{
				buddy->tree[index] = LMA_BUDDY_USED;
				char* mem = LMA_BUDDY_MEM_FROM_INDEX(buddy, index);
				
				/* loop back to set full flag */
				while(index > 1)
				{
					int left = ((index >> 1) << 1);
					int right = left + 1;
					if((buddy->tree[left] == LMA_BUDDY_USED
						|| buddy->tree[left] == LMA_BUDDY_FULL)
						&& (buddy->tree[right] == LMA_BUDDY_USED
						|| buddy->tree[right] == LMA_BUDDY_FULL))
					{
						index /= 2;
						buddy->tree[index] = LMA_BUDDY_FULL;
						continue;
					}	
					break;
				}

				/* return allocated memory */
				return mem;
			}

			/* split and going down to get required */
			else if(mem_size > malloc_size)
			{
				buddy->tree[index] = LMA_BUDDY_SPLIT;
				index *= 2;
				continue;
			}

			/* go back to parent's buddy */
			else
			{
				do { index /= 2; } 
				while((index & 1) && (index > 1));
				LMA_ALLOC_CHECK(index);
				index ++;
			}
		}

		else if(buddy->tree[index] == LMA_BUDDY_USED
			|| buddy->tree[index] == LMA_BUDDY_FULL)
		{
			/* try self or parent's buddy */
			while((index & 1) && (index > 1))
			{
				index /= 2;
			}
			LMA_ALLOC_CHECK(index);
			index ++;
		}

		else if(buddy->tree[index] == LMA_BUDDY_SPLIT)
		{
			size_t mem_size = LMA_BUDDY_SIZE_FROM_INDEX(buddy, index);
			if(mem_size <= malloc_size)
			{
				/* try self or parent's buddy */
				while((index & 1) && (index > 1))
				{
					index /= 2;
				}
				LMA_ALLOC_CHECK(index);
				index ++;
			}

			/* go to child */
			else
			{
				index *= 2;
			}
		}
	}
#endif 

	return NULL;
}

void buddy_free(struct buddy_t* buddy, void* mem)
{
	if(!mem || !buddy) 
	{
		return;
	}
	
#ifdef __x86_64__
	int offset = (int)((uint64_t)mem - (uint64_t)buddy->pool);
#else
	int offset = (int)((uint32_t)mem - (uint32_t)buddy->pool);
#endif
	assert(offset <= buddy->pool_size);

	int index = 1;
	int step = buddy->pool_size;
	int current = 0;
	while(1)
	{	
		if(buddy->tree[index] == LMA_BUDDY_USED)
		{
			assert(current == offset);
			buddy->tree[index] = LMA_BUDDY_UNUSED;

			/* try merge buddy & set parent unused */
			int loop = index;
			while(loop > 1)
			{
				int left = ((loop >> 1) << 1);
				int right = left + 1;
				if(buddy->tree[left] == LMA_BUDDY_UNUSED 
					&& buddy->tree[right] == LMA_BUDDY_UNUSED)
				{
					loop /= 2;
					buddy->tree[loop] = LMA_BUDDY_UNUSED;
					continue;
				}
				break;
			}

			/* set full parent to split status */
			loop = index;
			while(loop > 1)
			{
				loop /= 2;
				if(buddy->tree[loop] == LMA_BUDDY_FULL)
				{
					buddy->tree[loop] = LMA_BUDDY_SPLIT;
					continue;
				}
				break;
			}
			
			return;
		}

		/* find dest by dichotomy */
		step/= 2;
		index *= 2;
		if(offset >= current + step)
		{
			current += step;
			index ++;
		}
	}
}

void lma_buddy_debug_unit(struct buddy_t* buddy, int index)
{
	int shift, size;
	assert(buddy);
	switch(buddy->tree[index])
	{
		case LMA_BUDDY_UNUSED:
			shift = LMA_BUDDY_SHIFT_FROM_INDEX(buddy, index);
			size = LMA_BUDDY_SIZE_FROM_INDEX(buddy, index);
			printf("(%d:%d) ", shift, shift + size);
			break;

		case LMA_BUDDY_FULL:
			printf("{ ");
			lma_buddy_debug_unit(buddy, index * 2);
			lma_buddy_debug_unit(buddy, index * 2 + 1);
			printf("} ");
			break;
			
		case LMA_BUDDY_SPLIT:
			lma_buddy_debug_unit(buddy, index * 2);
			lma_buddy_debug_unit(buddy, index * 2 + 1);
			break;

		case LMA_BUDDY_USED:
			shift = LMA_BUDDY_SHIFT_FROM_INDEX(buddy, index);
			size = LMA_BUDDY_SIZE_FROM_INDEX(buddy, index);
			printf("[%d:%d] ", shift, shift + size);
			break;

		default:
			assert(0);
	}
}

void buddy_debug(struct buddy_t* buddy)
{
	if(!buddy)  return;
	lma_buddy_debug_unit(buddy, 1);
	printf("\n");
}


